'''
AVL 动态空间实现
'''
class Node:

    def __init__(self, key, value, left=None, right=None):
        self.key = key
        self.value = value
        self.left = left
        self.right = right
        self.height = 1
    def __repr__(self):
        return f'key: {self.key}, value: {self.value}, height: {self.height}'


class AVL:
    def __init__(self):
        self.root = None
        self.size = 0

    def _up(self, node):  # 修正信息
        node.height = 1 + max(node.left.height if node.left else 0, node.right.height if node.right else 0)

    def _left_rotate(self, node):  # 左旋,左边节点为新头，同时保证仍为搜索树
        right = node.right  # 不判空，因为只有不平衡才调用，所以默认右肯定有
        node.right = right.left
        right.left = node
        self._up(node)  # 先维护自己，自己沉下去了
        self._up(right)  # 再维护新头节点
        return right

    def _right_rotate(self, node):  # 与左旋对称相似
        left = node.left
        node.left = left.right
        left.right = node
        self._up(node)
        self._up(left)
        return left

    def _maintain(self, node):
        # 维护平衡[新头可能会变]。AVL有四种：
        #   LL(右旋), RR(左旋),
        #   LR(左节点的右节点为新头，左节点左旋后自己右旋), RL(与LR对称，右节点右旋，自己左旋)
        if node is None: return None
        lh, rh = node.left.height if node.left else 0, node.right.height if node.right else 0
        if lh - rh > 1:
            llh = node.left.left.height if node.left and node.left.left else 0
            lrh = node.left.right.height if node.left and node.left.right else 0
            if llh >= lrh:  # LL型（llh > lrh），或者当删除的时候同时导致了LL,LR (llh == lrh and lh - rh == 2)
                return self._right_rotate(node)
            else:  # LR 型
                node.left = self._left_rotate(node.left)
                return self._right_rotate(node)
        elif rh - lh > 1:
            rrh = node.right.right.height if node.right and node.right.right else 0
            rlh = node.right.left.height if node.right and node.right.left else 0
            # 对称
            if rrh >= rlh:
                return self._left_rotate(node)
            else:
                node.right = self._right_rotate(node.right)
                return self._left_rotate(node)
        return node  # 不违规，无需调整

    def _get(self, node, k):
        if node is None:
            return None
        elif k < node.key:
            return self._get(node.left, k)
        elif k > node.key:
            return self._get(node.right, k)
        return node

    def _add(self, node, k, v):
        if node is None:
            self.size += 1
            return Node(k, v)
        elif k < node.key:
            node.left = self._add(node.left, k, v)
        elif k > node.key:
            node.right = self._add(node.right, k, v)
        else:
            node.value = v
        self._up(node)
        return self._maintain(node)

    def _delete(self, node, k):
        if node is None:
            return None
        elif k < node.key:
            node.left = self._delete(node.left, k)
        elif k > node.key:
            node.right = self._delete(node.right, k)
        else:
            # 没左没右直接删，只有一个孩子节点，孩子节点替换我
            if not node.left and not node.right:
                node = None
            elif node.left and not node.right:
                node = node.left
            elif node.right and not node.left:
                node = node.right
            else:
                # 既有左又有右，选择右节点的最左节点替换我
                # 先删除它，同时保持平衡维护信息啥的，然后将它替换我
                most_left = node.right  # 右树的最左节点
                while most_left.left:
                    most_left = most_left.left
                node.right = self._delete_most_left(node.right, most_left)
                most_left.left = node.left
                most_left.right = node.right
                node = most_left
            self.size -= 1
        if node:
            self._up(node)
        return self._maintain(node)

    def _delete_most_left(self, node, most_left):
        if node == most_left: return most_left.right
        node.left = self._delete_most_left(node.left, most_left)
        self._up(node)
        return self._maintain(node)

    # ******************** 基本api ****************************
    def put(self, key, value):
        self.root = self._add(self.root, key, value)

    def get(self, key):
        target = self._get(self.root, key)
        if not target: raise KeyError(key)
        return target.value

    def remove(self, key):
        self.root = self._delete(self.root, key)

    def size(self):
        return self.size

    __getitem__ = get
    __setitem__ = put
    __delitem__ = remove
    __len__ = size

    def floor(self, key):
        def f(root):
            if root is None: return float('-inf')
            if key < root.key:
                return f(root.left)
            elif key > root.key:
                return max(root.key, f(root.right))
            else:
                return root.key

        return f(self.root)

    def ceil(self, key):
        def f(root):
            if root is None: return float('inf')
            if key < root.key:
                return min(root.key, f(root.left))
            elif key > root.key:
                return f(root.right)
            else:
                return root.key

        return f(self.root)
